import tensorflow as tf
from tensorflow.keras import applications
import time

def main():
    model = applications.ResNet50(weights=None)
    print("model = " + str(model))

    with tf.device('/device:XLA_DTU:0'):
        data = tf.random_uniform([16, 224, 224, 3])
        probs = model(data, training=True)
        opt = tf.train.GradientDescentOptimizer(0.01)
        target = tf.constant([777] * 16, dtype=tf.int32)
        loss = tf.losses.sparse_softmax_cross_entropy(target, probs)
        train_opt = opt.minimize(loss)
        init = tf.global_variables_initializer()

    config = tf.ConfigProto()
    config.log_device_placement = True

    with tf.compat.v1.Session(config=config) as sess:
        sess.run(init)
        sess.run(train_opt)


if __name__ == '__main__':
    main()